{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG Quickstart for Azure AI Search\n",
    "\n",
    "This quickstart provides a query for RAG scenarios. It demonstrates an approach for a chat experience using grounding data from a search index on Azure AI Search.\n",
    "\n",
    "We took a few shortcuts to keep the exercise basic and focused on query definitions:\n",
    "\n",
    "- We use the **hotels-sample-index**, which can be created in minutes and runs on any search service tier. This index is created by a wizard using built-in sample data.\n",
    "\n",
    "- We omit vectors so that we can skip chunking and embedding. The index contains plain text.\n",
    "\n",
    "A non-vector index isn't ideal for RAG patterns, but it makes for a simpler example.\n",
    "\n",
    "Once you understand the fundamentals of integrating queries from Azure AI Search to a chat completion model on Azure OpenAI, you can build on that experience by adding vector fields and vector and hybrid queries. We recommend the [phi-chat Python code example](https://github.com/Azure/azure-search-vector-samples/blob/main/demo-python/code/phi-chat/phi-chat.ipynb) for that step.\n",
    "\n",
    "This example is fully documented in [Quickstart: Generative search (RAG) with grounding data from Azure AI Search](https://learn.microsoft.com/azure/search/search-get-started-rag). If you need more guidance than the readme provides, please refer to the article.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "- [Azure OpenAI](https://learn.microsoft.com/azure/ai-services/openai/how-to/create-resource)\n",
    "\n",
    "  - [Choose a region](https://learn.microsoft.com/azure/ai-services/openai/concepts/models?tabs=global-standard%2Cstandard-chat-completions#global-standard-model-availability) that supports the chat completion model you want to use (gpt-4o, gpt-4o-mini, or equivalent model). \n",
    "  - [Deploy the chat completion model](https://learn.microsoft.com/azure/ai-studio/how-to/deploy-models-openai) in Azure AI Foundry or [use another approach](https://learn.microsoft.com/azure/ai-services/openai/how-to/working-with-models?tabs=powershell).\n",
    "\n",
    "\n",
    "- [Azure AI Search](https://learn.microsoft.com/azure/search/search-create-service-portal)\n",
    "\n",
    "  - Basic tier or higher is recommended.\n",
    "  - Enable semantic ranking.\n",
    "  - Enable role-based access control.\n",
    "  - Enable a system identity for Azure AI Search.\n",
    "  \n",
    "Make sure you know the name of the deployed model, and have the endpoints for both Azure resources at hand. You will provide this information in the steps that follow.\n",
    "\n",
    "## Configure access\n",
    "\n",
    "This quickstart assumes authentication and authorization using Microsoft Entra ID and role assignments. It also assumes that you run this code from your local device.\n",
    "\n",
    "1. To create, load, and query the sample hotels index on Azure AI Search, you must personally have role assignments for **Search Index Data Contributor** and **Search Service Contributor**. If the hotels index already exists, you only need **Search Index Data Reader**.\n",
    "\n",
    "1. To send the query and search results to Azure OpenAI, you must have **Cognitive Services OpenAI User** permissions on Azure OpenAI.\n",
    "\n",
    "## Create the sample index\n",
    "\n",
    "This quickstart assumes the hotels-sample-index, which you can create in minutes using [this quickstart](https://learn.microsoft.com/azure/search/search-get-started-portal).\n",
    "\n",
    "Once the index exists, modify it in the Azure portal to use this semantic configuration:"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "    \"semantic\": {\n",
    "    \"defaultConfiguration\": \"semantic-config\",\n",
    "    \"configurations\": [\n",
    "      {\n",
    "        \"name\": \"semantic-config\",\n",
    "        \"prioritizedFields\": {\n",
    "          \"titleField\": { \"fieldName\": \"HotelName\" },\n",
    "          \"prioritizedContentFields\": [ { \"fieldName\": \"Description\" } ],\n",
    "          \"prioritizedKeywordsFields\": [\n",
    "            { \"fieldName\": \"Category\" },\n",
    "            { \"fieldName\": \"Tags\" }\n",
    "          ]}\n",
    "       }]},"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have your Azure resources, an index, and model in place, you can run the script to chat with the index."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a virtual environment\n",
    "\n",
    "Create a virtual environment so that you can install the dependencies in isolation.\n",
    "\n",
    "1. In Visual Studio Code, open the folder containing Quickstart-RAG.ipynb.\n",
    "\n",
    "1. Press Ctrl-shift-P to open the command palette, search for \"Python: Create Environment\", and then select `Venv` to create a virtual environment in the current workspace.\n",
    "\n",
    "1. Select Quickstart-RAG\\requirements.txt for the dependencies.\n",
    "\n",
    "It takes several minutes to create the environment. When the environment is ready, continue to the next step.\n",
    "\n",
    "## Log in to Azure\n",
    "\n",
    "You're using Microsoft Entra ID and role assignments for the connection. Make sure you're logged in to the same tenant and subscription as Azure AI Search and Azure OpenAI. You can use the Azure CLI on the command line to show current properties, change properties, and to log in. For more information, see [Connect without keys](https://learn.microsoft.com/azure/search/search-get-started-rbac). \n",
    "\n",
    "Run each of the following commands in sequence.\n",
    "\n",
    "```\n",
    "az account show\n",
    "\n",
    "az account set --subscription <PUT YOUR SUBSCRIPTION ID HERE>\n",
    "\n",
    "az login --tenant <PUT YOUR TENANT ID HERE>\n",
    "```\n",
    "\n",
    "You should now be logged in to Azure from your local device.\n",
    "\n",
    "## Run the code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Package install for quickstart\n",
    "! pip install -r requirements.txt --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set endpoints and deployment model (provide the name of the deployment)\n",
    "AZURE_SEARCH_SERVICE: str = \"PUT YOUR SEARCH SERVICE ENDPOINT HERE\"\n",
    "AZURE_OPENAI_ACCOUNT: str = \"PUT YOUR AZURE OPENAI ENDPOINT HERE\"\n",
    "AZURE_DEPLOYMENT_MODEL: str = \"gpt-4o\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic RAG Query\n",
    "\n",
    "The following code demonstrates how to set up a query on string fields and string collections in a search index, and how to send the results to a chat completion model. The response returned to the user is from the chat completion model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the query for generating responses\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.identity import get_bearer_token_provider\n",
    "from azure.search.documents import SearchClient\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "credential = DefaultAzureCredential()\n",
    "token_provider = get_bearer_token_provider(credential, \"https://cognitiveservices.azure.com/.default\")\n",
    "openai_client = AzureOpenAI(\n",
    "    api_version=\"2024-06-01\",\n",
    "    azure_endpoint=AZURE_OPENAI_ACCOUNT,\n",
    "    azure_ad_token_provider=token_provider\n",
    ")\n",
    "\n",
    "search_client = SearchClient(\n",
    "    endpoint=AZURE_SEARCH_SERVICE,\n",
    "    index_name=\"hotels-sample-index\",\n",
    "    credential=credential\n",
    ")\n",
    "\n",
    "# This prompt provides instructions to the model\n",
    "GROUNDED_PROMPT=\"\"\"\n",
    "You are a friendly assistant that recommends hotels based on activities and amenities.\n",
    "Answer the query using only the sources provided below in a friendly and concise bulleted manner.\n",
    "Answer ONLY with the facts listed in the list of sources below.\n",
    "If there isn't enough information below, say you don't know.\n",
    "Do not generate answers that don't use the sources below.\n",
    "Query: {query}\n",
    "Sources:\\n{sources}\n",
    "\"\"\"\n",
    "\n",
    "# Query is the question being asked. It's sent to the search engine and the LLM.\n",
    "query=\"Can you recommend a few hotels with complimentary breakfast?\"\n",
    "\n",
    "# Search results are created by the search client.\n",
    "# Search results are composed of the top 5 results and the fields selected from the search index.\n",
    "# Search results include the top 5 matches to your query.\n",
    "search_results = search_client.search(\n",
    "    search_text=query,\n",
    "    top=5,\n",
    "    select=\"Description,HotelName,Tags\",\n",
    "    query_type=\"semantic\"\n",
    ")\n",
    "sources_formatted = \"\\n\".join([f'{document[\"HotelName\"]}:{document[\"Description\"]}:{document[\"Tags\"]}' for document in search_results])\n",
    "\n",
    "# Send the search results and the query to the LLM to generate a response based on the prompt.\n",
    "response = openai_client.chat.completions.create(\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": GROUNDED_PROMPT.format(query=query, sources=sources_formatted)\n",
    "        }\n",
    "    ],\n",
    "    model=AZURE_DEPLOYMENT_MODEL\n",
    ")\n",
    "\n",
    "# Here is the response from the chat model.\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you get an authorization error instead of results:\n",
    "\n",
    "+ If you just enabled role assignments, wait a few minutes and try again. It can take several minutes for role assignments to become operational.\n",
    "\n",
    "+ Make sure you [enabled RBAC](https://learn.microsoft.com/azure/search/search-security-enable-roles?tabs=config-svc-portal%2Cdisable-keys-portal) on Azure AI Search. An HttpStatusMessage of **Forbidden** is an indicator that RBAC isn't enabled.\n",
    "\n",
    "+ Recheck role assignments if you get a 401 error. You should also review the steps in [Connect without keys](https://learn.microsoft.com/azure/search/search-get-started-rbac).\n",
    "\n",
    "+ Check firewall settings. This quickstart assumes public network access. If you have a firewall, you need to add rules to allow inbound requests from your device and for service-to-service connections. For more information, see [Configure network access](https://learn.microsoft.com/azure/search/service-configure-firewall).\n",
    "\n",
    "+ For more debugging guidance, see the [troubleshooting section](https://learn.microsoft.com/azure/search/search-get-started-rag#troubleshooting-errors) in the Quickstart documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complex RAG Query\n",
    "\n",
    "Use complex types and collections in a RAG pattern by converting the result to JSON before sending it to the chat completion models. This query is similar to the basic query, but it includes fields from the \"Address\" complex type and the \"Rooms\" complex collection. This query shows you how to pull in data from parts of the index that aren't just simple text fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Query is the question being asked. It's sent to the search engine and the LLM.\n",
    "query=\"Can you recommend a few hotels that offer complimentary breakfast? Tell me their description, address, tags, and the rate for one room they have which sleep 4 people.\"\n",
    "\n",
    "# Set up the search results and the chat thread.\n",
    "# Retrieve the selected fields from the search index related to the question.\n",
    "selected_fields = [\"HotelName\",\"Description\",\"Address\",\"Rooms\",\"Tags\"]\n",
    "search_results = search_client.search(\n",
    "    search_text=query,\n",
    "    top=5,\n",
    "    select=selected_fields,\n",
    "    query_type=\"semantic\"\n",
    ")\n",
    "sources_filtered = [{field: result[field] for field in selected_fields} for result in search_results]\n",
    "sources_formatted = \"\\n\".join([json.dumps(source) for source in sources_filtered])\n",
    "\n",
    "response = openai_client.chat.completions.create(\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": GROUNDED_PROMPT.format(query=query, sources=sources_formatted)\n",
    "        }\n",
    "    ],\n",
    "    model=AZURE_DEPLOYMENT_MODEL\n",
    ")\n",
    "\n",
    "print(response.choices[0].message.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
